{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/pytorch/convert_pytorch_training_torchnano.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert PyTorch Training Loop to Use `TorchNano`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📚 **Related Reading**\n",
    ">\n",
    "> If you have already defined a PyTorch training loop function with a model, optimizers, and dataloaders as parameters, you could refer to [this guide](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/use_nano_decorator_pytorch_training.html) to use `@nano` decorator, which is a simpler way to gain acceleration from BigDL-Nano."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`TorchNano` API integrates multiple optimizations to accelerate custom PyTorch training loop. As a pure PyTorch user, you could apply few changes to your existing code to use `TorchNano`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "To use `TorchNano`, you need to install BigDL-Nano for PyTorch first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "!pip install --pre --upgrade bigdl-nano[pytorch] # install the nightly-built version\n",
    "!source bigdl-nano-init # set environment variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> Before starting your PyTorch application, it is highly recommended to run `source bigdl-nano-init` to set several environment variables based on your current hardware. Empirically, these variables will bring big performance increase for most PyTorch applications on training workloads."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "> ⚠️ **Warning**\n",
    "> \n",
    "> For Jupyter Notebook users, we recommend to run the commands above, especially `source bigdl-nano-init` before jupyter kernel is started, or some of the optimizations may not take effect."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PyTorch Training Loops Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose you would like to finetune a [ResNet-18 model](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html) (pretrained on ImageNet dataset) on [OxfordIIITPet dataset](https://pytorch.org/vision/main/generated/torchvision.datasets.OxfordIIITPet.html), you may create datasets, the model and define your training loops as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# Define model and dataloader\n",
    "\n",
    "from torch import nn\n",
    "from torchvision.models import resnet18\n",
    "\n",
    "class MyPytorchModule(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model = resnet18(pretrained=True)\n",
    "        num_ftrs = self.model.fc.in_features\n",
    "        # Here the size of each output sample is set to 37.\n",
    "        self.model.fc = nn.Linear(num_ftrs, 37)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import OxfordIIITPet\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "def create_train_dataloader():\n",
    "    train_transform = transforms.Compose([transforms.Resize(256),\n",
    "                                          transforms.RandomCrop(224),\n",
    "                                          transforms.RandomHorizontalFlip(),\n",
    "                                          transforms.ColorJitter(brightness=.5, hue=.3),\n",
    "                                          transforms.ToTensor(),\n",
    "                                          transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                                               [0.229, 0.224, 0.225])])\n",
    "\n",
    "    # apply data augmentation to the tarin_dataset\n",
    "    train_dataset = OxfordIIITPet(root=\"/tmp/data\", transform=train_transform, download=True)\n",
    "\n",
    "    # prepare data loaders\n",
    "    train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "\n",
    "    return train_dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "def train_loops():\n",
    "    model = MyPytorchModule()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "    loss_fuc = torch.nn.CrossEntropyLoss()\n",
    "    train_loader = create_train_dataloader()\n",
    "\n",
    "    num_epochs = 5\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "\n",
    "        model.train()\n",
    "        train_loss, num = 0, 0\n",
    "        with tqdm(train_loader, unit=\"batch\") as tepoch:\n",
    "            for data, target in tepoch:\n",
    "                tepoch.set_description(f\"Epoch {epoch}\")\n",
    "                optimizer.zero_grad()\n",
    "                output = model(data)\n",
    "                loss = loss_fuc(output, target)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                loss_value = loss.sum()\n",
    "                train_loss += loss_value\n",
    "                num += 1\n",
    "                tepoch.set_postfix(loss=loss_value)\n",
    "        print(f'Train Epoch: {epoch}, avg_loss: {train_loss / num}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; _The definition of_ `MyPytorchModule` _and_ `create_train_dataloader` _can be found in the_ [runnable example](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/pytorch/convert_pytorch_training_torchnano.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to `TorchNano`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are 5 simple steps to convert your PyTorch code to use `TorchNano`:\n",
    "\n",
    "1. Import `TorchNano`\n",
    "2. Subclass `TorchNano` and override its `train` method\n",
    "3. Move the code for your custom training loops inside the `TorchNano`'s `train` method\n",
    "4. Call `TorchNano`'s `setup` method to set up model, optimizer(s), and dataloader(s) for accelerated training\n",
    "5. Replace `loss.backward()` with `self.backward(loss)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1. import TorchNano\n",
    "from bigdl.nano.pytorch import TorchNano\n",
    "\n",
    "# Step 2. subclass TorchNano and override its train method\n",
    "class MyNano(TorchNano):\n",
    "    def train(self):\n",
    "        # Step 3. Move the code for your custom training loops inside the train method\n",
    "        model = MyPytorchModule()\n",
    "        optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "        loss_fuc = torch.nn.CrossEntropyLoss()\n",
    "        train_loader = create_train_dataloader()\n",
    "\n",
    "        # Step 4. call setup method to set up model, optimizer(s),\n",
    "        #         and dataloader(s) for accelerated training\n",
    "        model, optimizer, train_loader = self.setup(model, optimizer, train_loader)\n",
    "        num_epochs = 5\n",
    "\n",
    "        for epoch in range(num_epochs):\n",
    "\n",
    "            model.train()\n",
    "            train_loss, num = 0, 0\n",
    "            with tqdm(train_loader, unit=\"batch\") as tepoch:\n",
    "                for data, target in tepoch:\n",
    "                    tepoch.set_description(f\"Epoch {epoch}\")\n",
    "                    optimizer.zero_grad()\n",
    "                    output = model(data)\n",
    "                    loss = loss_fuc(output, target)\n",
    "                    # Step 5. Replace loss.backward() with self.backward(loss)\n",
    "                    self.backward(loss)\n",
    "                    optimizer.step()\n",
    "                    loss_value = loss.sum()\n",
    "                    train_loss += loss_value\n",
    "                    num += 1\n",
    "                    tepoch.set_postfix(loss=loss_value)\n",
    "            print(f'Train Epoch: {epoch}, avg_loss: {train_loss / num}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> To make sure that the converted `TorchNano` still has a functional training loop, there are some requirements:\n",
    ">\n",
    "> - there should be one and only one instance of `torch.nn.Module` as model in the training loop \n",
    "> - there should be at least one instance of `torch.optim.Optimizer` as optimizer in the training loop\n",
    "> - there should be at least one instance of `torch.utils.data.DataLoader` as dataloader in the training loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could then do the training by instantiating `MyNano` and calling its `train` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MyNano().train()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> Due to the optimized environment variables set by `source bigdl-nano-init`, you could already experience some training acceleration after converting your PyTorch code to use `TorchNano`.\n",
    "> \n",
    "> For more optimizations provided by `TorchNano`, you can refer to the Related Readings."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📚 **Related Readings**\n",
    "> \n",
    "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)\n",
    "> - [How to accelerate a PyTorch application on training workloads through Intel® Extension for PyTorch*](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/accelerate_pytorch_training_ipex.html)\n",
    "> - [How to accelerate a PyTorch application on training workloads through multiple instances](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/accelerate_pytorch_training_multi_instance.html)\n",
    "> - [How to use the channels last memory format in your PyTorch application for training](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/pytorch_training_channels_last.html)\n",
    "> - [How to conduct BFloat16 Mixed Precision training in your PyTorch application](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Training/PyTorch/accelerate_pytorch_training_bf16.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nano-pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.15 (default, Nov 24 2022, 21:12:53) \n[GCC 11.2.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "09344c7f3239fd422839751f876786d6b1a624c40f19af1b43cb2737f421c2b2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
